import socket
import wave
import time
import os
import struct
import numpy as np  # 需要 numpy 来处理音量 (pip install numpy)

# ================= CONFIGURATION =================
PICO_IP = "192.168.43.143"  # <--- 你的 Pico IP
PORT = 50005
# 建议加大包长到 1024，减少 UDP 包头开销，传输更稳
CHUNK_SIZE = 1024  

WAV_FILE = "ai_reply.wav"

# === 电脑端音量控制 ===
# 0.5 = 50% 音量 (Pico 直通声音很大，建议在此处衰减)
VOLUME = 0.5 
# =======================================

def play_audio_smart():
    if not os.path.exists(WAV_FILE):
        print(f"File {WAV_FILE} not found!")
        return

    wf = wave.open(WAV_FILE, 'rb')
    
    # 检查格式
    if wf.getnchannels() != 1 or wf.getframerate() != 16000 or wf.getsampwidth() != 2:
        print("Error: Format mismatch. Require: 16kHz, Mono, 16-bit")
        wf.close()
        return

    print(f"Loading and processing audio...")
    
    # 1. 一次性读取所有数据到内存 (处理音量更方便)
    all_frames = wf.readframes(wf.getnframes())
    wf.close()
    
    # 2. 转换成 numpy 数组进行音量缩放
    audio_data = np.frombuffer(all_frames, dtype=np.int16)
    # 乘上音量系数，防止炸麦
    audio_data = (audio_data * VOLUME).astype(np.int16)
    # 转回字节流
    payload = audio_data.tobytes()
    
    total_bytes = len(payload)
    print(f"Total size: {total_bytes} bytes. Sending to {PICO_IP}:{PORT} ...")

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    
    # 3. 智能发送循环
    # 理论传输速率: 32000 bytes/sec
    BYTES_PER_SEC = 32000
    
    offset = 0
    # 获取高精度启动时间
    start_time = time.perf_counter()
    
    try:
        while offset < total_bytes:
            # 取出一个切片
            end = min(offset + CHUNK_SIZE, total_bytes)
            chunk = payload[offset:end]
            
            # 发送
            sock.sendto(chunk, (PICO_IP, PORT))
            offset += len(chunk)
            
            # === [核心算法] 自适应流控 ===
            # 计算理论上此时应该发送了多少时间
            expected_time = offset / BYTES_PER_SEC
            
            # 计算实际上过去了多少时间
            actual_time = time.perf_counter() - start_time
            
            # 只有当我们发太快了才睡觉
            # 如果 actual_time > expected_time (发慢了)，sleep_time 会是负数，
            # 程序就不会休眠，全速发送直到追上进度。
            sleep_time = expected_time - actual_time
            
            if sleep_time > 0:
                time.sleep(sleep_time)
            
            # 简单的进度条
            if offset % (CHUNK_SIZE * 10) == 0:
                print(">", end="", flush=True)

        print("\nTransmission complete!")
        sock.sendto(b'END_OF_STREAM', (PICO_IP, PORT))

    except KeyboardInterrupt:
        print("\nStopped by user")
    finally:
        sock.close()

if __name__ == "__main__":
    play_audio_smart()